package org.apache.shardingsphere.transaction.base.seata.at;

import com.google.common.base.Preconditions;
import io.seata.config.FileConfiguration;
import io.seata.core.context.RootContext;
import io.seata.core.exception.TransactionException;
import io.seata.core.rpc.netty.RmNettyRemotingClient;
import io.seata.core.rpc.netty.TmNettyRemotingClient;
import io.seata.rm.RMClient;
import io.seata.rm.datasource.DataSourceProxy;
import io.seata.tm.TMClient;
import io.seata.tm.api.GlobalTransaction;
import io.seata.tm.api.GlobalTransactionContext;
import lombok.SneakyThrows;
import org.apache.shardingsphere.infra.database.type.DatabaseType;
import org.apache.shardingsphere.transaction.core.ResourceDataSource;
import org.apache.shardingsphere.transaction.core.TransactionType;
import org.apache.shardingsphere.transaction.spi.ShardingSphereTransactionManager;

import javax.sql.DataSource;
import java.sql.Connection;
import java.sql.SQLException;
import java.util.Collection;
import java.util.HashMap;
import java.util.Map;

/**
 * Seata AT ShardingSphere 事务管理器
 *
 * @Author Hsu琛君珩
 * @Date 2024-10-05 5:00
 * @Description
 * @Version: v1.0.0
 */
@SuppressWarnings("AlibabaClassNamingShouldBeCamel")
public class SeataATShardingSphereTransactionManager implements ShardingSphereTransactionManager {

    // 存储数据源映射
    private final Map<String, DataSource> dataSourceMap = new HashMap<>();

    // 应用ID
    private final String applicationId;

    // 事务服务组
    private final String transactionServiceGroup;

    // 是否启用 Seata AT
    private final boolean enableSeataAT;

    // 全局事务超时时间
    private final int globalTXTimeout;

    /**
     * 构造方法，从配置文件加载Seata相关配置。
     */
    public SeataATShardingSphereTransactionManager() {
        FileConfiguration config = new FileConfiguration("seata.conf");
        enableSeataAT = config.getBoolean("sharding.transaction.seata.at.enable", true);
        applicationId = config.getConfig("client.application.id");
        transactionServiceGroup = config.getConfig("client.transaction.service.group", "default");
        globalTXTimeout = config.getInt("sharding.transaction.seata.tx.timeout", 60);
    }

    /**
     * 初始化 Seata RPC 客户端和数据源代理。
     *
     * @param databaseType 数据库类型
     * @param resourceDataSources 数据源集合
     * @param providerType 提供者类型
     */
    @Override
    public void init(final DatabaseType databaseType, final Collection<ResourceDataSource> resourceDataSources, final String providerType) {
        if (enableSeataAT) {
            initSeataRPCClient();  // 初始化 Seata RPC 客户端
            resourceDataSources.forEach(each -> dataSourceMap.put(each.getOriginalName(), new DataSourceProxy(each.getDataSource())));
        }
    }

    /**
     * 初始化 Seata TM 和 RM 客户端。
     */
    private void initSeataRPCClient() {
        Preconditions.checkNotNull(applicationId, "请在 seata.conf 文件中配置 application id.");
        TMClient.init(applicationId, transactionServiceGroup);  // 初始化 TM 客户端
        RMClient.init(applicationId, transactionServiceGroup);  // 初始化 RM 客户端
    }

    /**
     * 获取事务类型。
     *
     * @return 事务类型，Seata 基础事务
     */
    @Override
    public TransactionType getTransactionType() {
        return TransactionType.BASE;
    }

    /**
     * 判断当前是否在事务中。
     *
     * @return 是否在事务中
     */
    @Override
    public boolean isInTransaction() {
        Preconditions.checkState(enableSeataAT, "sharding seata-at transaction has been disabled.");
        return null != SeataTransactionHolder.get();  // 检查事务上下文是否存在
    }

    /**
     * 获取数据库连接。
     *
     * @param databaseName 数据库名称
     * @param dataSourceName 数据源名称
     * @return 数据库连接
     * @throws SQLException 发生SQL异常时抛出
     */
    @Override
    public Connection getConnection(final String databaseName, final String dataSourceName) throws SQLException {
        Preconditions.checkState(enableSeataAT, "sharding seata-at transaction has been disabled.");
        return dataSourceMap.get(databaseName + "." + dataSourceName).getConnection();  // 获取分表数据源的连接
    }

    /**
     * 开始全局事务。
     */
    @Override
    public void begin() {
        begin(globalTXTimeout);  // 使用默认超时时间开始事务
    }

    /**
     * 开始全局事务，指定超时时间。
     *
     * @param timeout 超时时间（秒）
     * @throws TransactionException 当事务开启失败时抛出
     */
    @Override
    @SneakyThrows(TransactionException.class)
    public void begin(final int timeout) {
        if (timeout < 0) {
            throw new TransactionException("超时时间应大于0秒");
        }
        Preconditions.checkState(enableSeataAT, "sharding seata-at transaction has been disabled.");
        GlobalTransaction globalTransaction = GlobalTransactionContext.getCurrentOrCreate();  // 创建或获取全局事务
        globalTransaction.begin(timeout * 1000);  // 开始全局事务
        SeataTransactionHolder.set(globalTransaction);  // 设置事务上下文
    }

    /**
     * 提交全局事务。
     *
     * @param rollbackOnly 是否只回滚
     * @throws TransactionException 提交事务时异常
     */
    @Override
    @SneakyThrows(TransactionException.class)
    public void commit(final boolean rollbackOnly) {
        Preconditions.checkState(enableSeataAT, "sharding seata-at transaction has been disabled.");
        try {
            SeataTransactionHolder.get().commit();  // 提交全局事务
        } finally {
            SeataTransactionHolder.clear();  // 清除事务上下文
            RootContext.unbind();  // 解除RootContext绑定
        }
    }

    /**
     * 回滚全局事务。
     *
     * @throws TransactionException 回滚事务时异常
     */
    @Override
    @SneakyThrows(TransactionException.class)
    public void rollback() {
        Preconditions.checkState(enableSeataAT, "sharding seata-at transaction has been disabled.");
        try {
            SeataTransactionHolder.get().rollback();  // 回滚全局事务
        } finally {
            SeataTransactionHolder.clear();  // 清除事务上下文
            RootContext.unbind();  // 解除RootContext绑定
        }
    }

    /**
     * 关闭事务管理器，清理资源。
     */
    @Override
    public void close() {
        dataSourceMap.clear();  // 清空数据源映射
        SeataTransactionHolder.clear();  // 清除事务上下文
        RmNettyRemotingClient.getInstance().destroy();  // 销毁RM客户端
        TmNettyRemotingClient.getInstance().destroy();  // 销毁TM客户端
    }
}
